{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "## Rerun imports and initialization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import annotations\n",
    "\n",
    "import math\n",
    "import uuid\n",
    "from collections import namedtuple\n",
    "from math import cos, sin\n",
    "\n",
    "import numpy as np\n",
    "import rerun as rr  # pip install rerun-sdk\n",
    "import rerun.blueprint as rrb\n",
    "from rerun.notebook import Viewer  # pip install rerun-notebook"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "2",
   "metadata": {},
   "source": [
    "## Helper to create the colored cube\n",
    "\n",
    "This code exists in the `rerun.utilities` package, but is included here for context."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "ColorGrid = namedtuple(\"ColorGrid\", [\"positions\", \"colors\"])\n",
    "\n",
    "\n",
    "def build_color_grid(x_count: int = 10, y_count: int = 10, z_count: int = 10, twist: float = 0) -> ColorGrid:\n",
    "    \"\"\"\n",
    "    Create a cube of points with colors.\n",
    "\n",
    "    The total point cloud will have x_count * y_count * z_count points.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    x_count, y_count, z_count:\n",
    "        Number of points in each dimension.\n",
    "    twist:\n",
    "        Angle to twist from bottom to top of the cube\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    grid = np.mgrid[\n",
    "        slice(-x_count, x_count, x_count * 1j),\n",
    "        slice(-y_count, y_count, y_count * 1j),\n",
    "        slice(-z_count, z_count, z_count * 1j),\n",
    "    ]\n",
    "\n",
    "    angle = np.linspace(-float(twist) / 2, float(twist) / 2, z_count)\n",
    "    for z in range(z_count):\n",
    "        xv, yv, zv = grid[:, :, :, z]\n",
    "        rot_xv = xv * cos(angle[z]) - yv * sin(angle[z])\n",
    "        rot_yv = xv * sin(angle[z]) + yv * cos(angle[z])\n",
    "        grid[:, :, :, z] = [rot_xv, rot_yv, zv]\n",
    "\n",
    "    positions = np.vstack([xyz.ravel() for xyz in grid])\n",
    "\n",
    "    colors = np.vstack([\n",
    "        xyz.ravel()\n",
    "        for xyz in np.mgrid[\n",
    "            slice(0, 255, x_count * 1j),\n",
    "            slice(0, 255, y_count * 1j),\n",
    "            slice(0, 255, z_count * 1j),\n",
    "        ]\n",
    "    ])\n",
    "\n",
    "    return ColorGrid(positions.T, colors.T.astype(np.uint8))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4",
   "metadata": {},
   "source": [
    "## Logging some data\n",
    "\n",
    "Now we can log some data and add it to the recording, and display it using `notebook_show`.\n",
    "\n",
    "Note that displaying a recording will consume the data, so it will not be available for use in a subsequent cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "rr.init(\"rerun_example_cube\")\n",
    "\n",
    "STEPS = 100\n",
    "twists = math.pi * np.sin(np.linspace(0, math.tau, STEPS)) / 4\n",
    "for t in range(STEPS):\n",
    "    rr.set_time(\"step\", sequence=t)\n",
    "    cube = build_color_grid(10, 10, 10, twist=twists[t])\n",
    "    rr.log(\"cube\", rr.Points3D(cube.positions, colors=cube.colors, radii=0.5))\n",
    "\n",
    "rr.notebook_show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "## Logging live data\n",
    "\n",
    "Using `rr.notebook_show` like above buffers the data in the recording stream, but doesn't process it until the call to `rr.notebook_show`.\n",
    "\n",
    "However, `rr.notebook_show` can be called at any time during your cell's execution to immediately display the Rerun viewer. You can then incrementally stream to it. Here we add a sleep to simulate a cell that does a lot more processing. By calling `notebook_show` first, we can see the output of our code live while it's still running."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from time import sleep\n",
    "\n",
    "rr.init(\"rerun_example_cube\")\n",
    "\n",
    "rr.notebook_show()\n",
    "\n",
    "STEPS = 100\n",
    "twists = math.pi * np.sin(np.linspace(0, math.tau, STEPS)) / 4\n",
    "for t in range(STEPS):\n",
    "    sleep(0.05)\n",
    "    rr.set_time(\"step\", sequence=t)\n",
    "    cube = build_color_grid(10, 10, 10, twist=twists[t])\n",
    "    rr.log(\"cube\", rr.Points3D(cube.positions, colors=cube.colors, radii=0.5))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "metadata": {},
   "source": [
    "## Incremental logging\n",
    "\n",
    "Note that until we either reset the recording stream (by calling `rr.init()`), or create a new output widget (by calling `rr.notebook_show()` The currently active stream in the kernel will continue to send events to the existing widget.\n",
    "\n",
    "The following will add a rotation to the above recording."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "for t in range(STEPS):\n",
    "    sleep(0.01)\n",
    "    rr.set_time(\"step\", sequence=t)\n",
    "    rr.log(\"cube\", rr.Transform3D(rotation=rr.RotationAxisAngle(axis=[1, 0, 0], degrees=t)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10",
   "metadata": {},
   "source": [
    "## Starting a new recording\n",
    "\n",
    "You can always start another recording by calling `rr.init(...)` again to reset the global stream, or alternatively creating a separate recording stream using the `rr.RecordingStream` constructor (discussed more below)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "rr.init(\"rerun_example_cube\")\n",
    "\n",
    "STEPS = 100\n",
    "twists = math.pi * np.sin(np.linspace(0, math.tau, STEPS)) / 4\n",
    "for t in range(STEPS):\n",
    "    rr.set_time(\"step\", sequence=t)\n",
    "    h_grid = build_color_grid(10, 3, 3, twist=twists[t])\n",
    "    rr.log(\"h_grid\", rr.Points3D(h_grid.positions, colors=h_grid.colors, radii=0.5))\n",
    "    v_grid = build_color_grid(3, 3, 10, twist=twists[t])\n",
    "    rr.log(\"v_grid\", rr.Points3D(v_grid.positions, colors=v_grid.colors, radii=0.5))\n",
    "\n",
    "rr.notebook_show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12",
   "metadata": {},
   "source": [
    "## Adjusting the view\n",
    "\n",
    "The  `notebook_show` method also lets you adjust properties such as width and height."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "rr.init(\"rerun_example_cube\")\n",
    "\n",
    "small_cube = build_color_grid(3, 3, 3, twist=0)\n",
    "rr.log(\"small_cube\", rr.Points3D(small_cube.positions, colors=small_cube.colors, radii=0.5))\n",
    "\n",
    "rr.notebook_show(height=400)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14",
   "metadata": {},
   "source": [
    "To update the default width and height, you can use the `rerun.notebook.set_default_size` function.\n",
    "\n",
    "Note that you do not need to initialize a recording to use this function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rerun.notebook import set_default_size\n",
    "\n",
    "set_default_size(width=400, height=400)\n",
    "\n",
    "rr.init(\"rerun_example_cube\")\n",
    "\n",
    "small_cube = build_color_grid(3, 3, 3, twist=0)\n",
    "rr.log(\"small_cube\", rr.Points3D(small_cube.positions, colors=small_cube.colors, radii=0.5))\n",
    "\n",
    "rr.notebook_show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "set_default_size(width=640, height=480)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "17",
   "metadata": {},
   "source": [
    "## Using blueprints\n",
    "\n",
    "Rerun blueprints can be used with `rr.notebook_show()`\n",
    "\n",
    "For example, we can split the two grids into their own respective views."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [],
   "source": [
    "rr.init(\"rerun_example_cube\")\n",
    "\n",
    "blueprint = rrb.Blueprint(\n",
    "    rrb.Horizontal(\n",
    "        rrb.Spatial3DView(name=\"Horizontal grid\", origin=\"h_grid\"),\n",
    "        rrb.Spatial3DView(name=\"Vertical grid\", origin=\"v_grid\"),\n",
    "        column_shares=[2, 1],\n",
    "    ),\n",
    "    collapse_panels=True,\n",
    ")\n",
    "\n",
    "rr.notebook_show(blueprint=blueprint)\n",
    "\n",
    "STEPS = 100\n",
    "twists = math.pi * np.sin(np.linspace(0, math.tau, STEPS)) / 4\n",
    "for t in range(STEPS):\n",
    "    rr.set_time(\"step\", sequence=t)\n",
    "    h_grid = build_color_grid(10, 3, 3, twist=twists[t])\n",
    "    rr.log(\"h_grid\", rr.Points3D(h_grid.positions, colors=h_grid.colors, radii=0.5))\n",
    "    v_grid = build_color_grid(3, 3, 10, twist=twists[t])\n",
    "    rr.log(\"v_grid\", rr.Points3D(v_grid.positions, colors=v_grid.colors, radii=0.5))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19",
   "metadata": {},
   "source": [
    "## Extra convenience\n",
    "\n",
    "Rerun blueprints types also implement `_ipython_display_()` directly, so if a blueprint is the last element in your cell the right thing will happen.\n",
    "\n",
    "Note that this mechanism only works when you are using the global recording stream."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20",
   "metadata": {},
   "outputs": [],
   "source": [
    "rr.init(\"rerun_example_cube\")\n",
    "\n",
    "STEPS = 100\n",
    "twists = math.pi * np.sin(np.linspace(0, math.tau, STEPS)) / 4\n",
    "for t in range(STEPS):\n",
    "    rr.set_time(\"step\", sequence=t)\n",
    "    h_grid = build_color_grid(10, 3, 3, twist=twists[t])\n",
    "    rr.log(\"h_grid\", rr.Points3D(h_grid.positions, colors=h_grid.colors, radii=0.5))\n",
    "    v_grid = build_color_grid(3, 3, 10, twist=twists[t])\n",
    "    rr.log(\"v_grid\", rr.Points3D(v_grid.positions, colors=v_grid.colors, radii=0.5))\n",
    "\n",
    "rrb.Spatial3DView(name=\"Horizontal grid\", origin=\"h_grid\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21",
   "metadata": {},
   "source": [
    "## Working with non-global streams\n",
    "\n",
    "Sometimes it can be more explicit to work with specific (non-global recording) streams via `rr.RecordingStream` constructor.\n",
    "\n",
    "In this case, remember to call `notebook_show` directly on the recording stream. As noted above, there is no way to use a bare Blueprint object in conjunction with a non-global recording."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22",
   "metadata": {},
   "outputs": [],
   "source": [
    "rec = rr.RecordingStream(\"rerun_example_cube_flat\")\n",
    "\n",
    "bp = rrb.Blueprint(collapse_panels=True)\n",
    "\n",
    "rec.notebook_show(blueprint=bp)\n",
    "\n",
    "flat_grid = build_color_grid(20, 20, 1, twist=0)\n",
    "rec.log(\"flat_grid\", rr.Points3D(flat_grid.positions, colors=flat_grid.colors, radii=0.5))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23",
   "metadata": {},
   "source": [
    "## Using the Viewer object directly\n",
    "\n",
    "Instead of calling `notebook_show` you can alternatively hold onto the viewer object.\n",
    "\n",
    "This lets you add additional recordings to a view."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24",
   "metadata": {},
   "outputs": [],
   "source": [
    "rec = rr.RecordingStream(\"rerun_example_multi_recording\", recording_id=uuid.uuid4())\n",
    "\n",
    "flat_grid = build_color_grid(20, 20, 1, twist=0)\n",
    "rec.log(\"flat_grid\", rr.Points3D(flat_grid.positions, colors=flat_grid.colors, radii=0.5))\n",
    "\n",
    "viewer = Viewer(recording=rec, blueprint=rrb.Blueprint(rrb.BlueprintPanel(state=\"expanded\")))\n",
    "viewer.display()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25",
   "metadata": {},
   "outputs": [],
   "source": [
    "rec = rr.RecordingStream(\"rerun_example_multi_recording\", recording_id=uuid.uuid4())\n",
    "\n",
    "viewer.add_recording(rec)\n",
    "\n",
    "STEPS = 100\n",
    "twists = math.pi * np.sin(np.linspace(0, math.tau, STEPS)) / 4\n",
    "for t in range(STEPS):\n",
    "    rec.set_time(\"step\", sequence=t)\n",
    "    cube = build_color_grid(10, 10, 10, twist=twists[t])\n",
    "    rec.log(\"cube\", rr.Points3D(cube.positions, colors=cube.colors, radii=0.5))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26",
   "metadata": {},
   "source": [
    "## Controlling the Viewer\n",
    "\n",
    "Other than sending a blueprint to the Viewer, some parts of it can also be controlled directly through Python."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27",
   "metadata": {},
   "outputs": [],
   "source": [
    "viewer = Viewer()\n",
    "viewer.display()\n",
    "\n",
    "recordings = [\n",
    "    rr.RecordingStream(\"rerun_example_time_ctrl\", recording_id=\"example_a\"),\n",
    "    rr.RecordingStream(\"rerun_example_time_ctrl\", recording_id=\"example_b\"),\n",
    "]\n",
    "\n",
    "rec_colors = {\"example_a\": [0, 255, 0], \"example_b\": [255, 0, 0]}\n",
    "\n",
    "for rec in recordings:\n",
    "    viewer.add_recording(rec)\n",
    "\n",
    "STEPS = 100\n",
    "twists = math.pi * np.sin(np.linspace(0, math.tau, STEPS)) / 4\n",
    "for rec in recordings:\n",
    "    for t in range(STEPS):\n",
    "        cube = build_color_grid(10, 10, 10, twist=twists[t])\n",
    "        rec.set_time(\"step\", sequence=t)\n",
    "        rec.log(\"cube\", rr.Points3D(cube.positions, colors=rec_colors[rec.get_recording_id()], radii=0.5))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28",
   "metadata": {},
   "source": [
    "The state of each panel in the viewer can be overridden, locking it in the specified state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29",
   "metadata": {},
   "outputs": [],
   "source": [
    "viewer.update_panels(blueprint=\"expanded\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30",
   "metadata": {},
   "source": [
    "In multi-recording scenarios, the active recording can be set using `set_active_recording`. The timeline panel's state for the currently active recording can be controlled using `set_time_ctrl`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31",
   "metadata": {},
   "outputs": [],
   "source": [
    "viewer.set_active_recording(recording_id=\"example_a\")\n",
    "viewer.set_time_ctrl(timeline=\"step\", sequence=25)\n",
    "viewer.set_active_recording(recording_id=\"example_b\")\n",
    "viewer.set_time_ctrl(timeline=\"step\", sequence=75)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32",
   "metadata": {},
   "source": [
    "Switch between the two recordings in the blueprint panel to see the updated timelines."
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 5
}
